package resolver

import (
	"context"
	"go-rpc-2/server_register_and_found/register"
	"google.golang.org/grpc/attributes"
	"google.golang.org/grpc/resolver"
	"time"
)

type Builder struct {
	Registry register.Register
	duration time.Duration
}

func (b *Builder) Build(target resolver.Target, cc resolver.ClientConn, opts resolver.BuildOptions) (resolver.Resolver, error) {
	r := &Resolver{
		cc:       cc,
		target:   target,
		register: b.Registry,
		duration: b.duration,
		close:    make(chan struct{}),
	}
	r.ResolveNow(resolver.ResolveNowOptions{}) // 需要立即调用一次
	go r.watch()
	return r, nil
}

func (b *Builder) Scheme() string {
	return "dns"
}

func NewBuilder(register register.Register, duration time.Duration) *Builder {
	return &Builder{
		Registry: register,
		duration: duration,
	}
}

type Resolver struct {
	cc       resolver.ClientConn
	target   resolver.Target
	register register.Register
	duration time.Duration
	close    chan struct{}
}

func (r *Resolver) watch() {
	event := r.register.Subscribe(r.target.Endpoint())
	for {
		select {
		case <-r.close:
			return
		case <-event:
			r.resolve(resolver.ResolveNowOptions{})
		}
	}
}

func (r *Resolver) resolve(options resolver.ResolveNowOptions) {
	ctx, cancel := context.WithTimeout(context.Background(), r.duration)
	si, err := r.register.ListService(ctx, r.target.Endpoint())
	cancel()
	if err != nil {
		r.cc.ReportError(err)
	}
	Addresses := make([]resolver.Address, len(si))
	for i, s := range si {
		Addresses[i] = resolver.Address{
			Addr:       s.Addr,
			Attributes: attributes.New("weight", s.Weight).WithValue("group", s.Group),
		}
	}

	err = r.cc.UpdateState(resolver.State{
		Addresses: Addresses,
	})
	if err != nil {
		r.cc.ReportError(err)
	}
}

func (r *Resolver) ResolveNow(options resolver.ResolveNowOptions) {
	r.resolve(options)
}

func (r *Resolver) Close() {
	close(r.close)
}
